
from torch import  nn ,optim
from torch.utils.data import DataLoader
import torch
from date_set import *
from yolo_net import *
from torchsummary import summary
from configure import *

def loss_fun(output,target,no_obj,xishu):

    output = output.permute(0,2,3,1)
    output = output.reshape(output.size(0),output.size(1),output.size(2),3,-1)

    
    loss_obj_fun = nn.BCELoss()
    loss_obj = torch.mean(loss_obj_fun(torch.sigmoid(output[...,0]),target[...,0])*target[...,0])

    loss_noobj_fun = nn.BCELoss()
    loss_noobj = torch.mean(loss_noobj_fun(torch.sigmoid(output[...,0]),target[...,0])*no_obj[...,0])

    loss_conf = loss_obj + 0.3*loss_noobj

    loss_xy = nn.BCELoss()

    loss_box_x =  torch.mean(loss_xy(torch.sigmoid(output[...,1]),target[...,1])*target[...,0])
    loss_box_y =  torch.mean(loss_xy(torch.sigmoid(output[...,2]),target[...,2])*target[...,0])
    loss_box_fun = nn.MSELoss()
    
    loss_box_w =  torch.mean(loss_box_fun(output[...,3],target[...,3])*target[...,0]*xishu[...,0])
    loss_box_h =  torch.mean(loss_box_fun(output[...,4],target[...,4])*target[...,0]*xishu[...,0])


    loss_class_fun = nn.BCELoss()
    loss_class = torch.mean(loss_class_fun(torch.sigmoid(output[...,5:]),target[...,5:])*target[...,0])

    return loss_box_x+loss_box_y+loss_box_w+ loss_box_h+ loss_class+ loss_conf

if __name__ =='__main__':
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') 
    date_set = Yolo_dataset()
    
    net = Yolo_all().to(device=device)
    
    date_load = DataLoader(dataset=date_set,batch_size=6,shuffle=True)
    weigh_path = 'E:/python file/yolo_advance/weight/yolo_net.pt'
    if os.path.exists(weigh_path):
        net.load_state_dict(torch.load(weigh_path))
        print('*'*30+'加载权重成功'+'*'*30)
    opt = optim.Adam(params= net.parameters(),lr=0.001)
    epoch =0
    history_best = np.float16('inf')
    while epoch<500:
        for para in opt.param_groups:
            para['lr'] *= 0.99
        for image,traget13,target26,traget52,traget_no_13,traget_no_26,traget_no_52,xishu_13, xishu_26,xishu_52 in date_load:
            image,target13,target26,target52,traget_no_13,traget_no_26,traget_no_52,xishu_13, xishu_26,xishu_52 = image.to(device), \
                traget13.to(device),target26.to(device),traget52.to(device),traget_no_13.to(device),traget_no_26.to(device),traget_no_52.to(device), \
                xishu_13.to(device), xishu_26.to(device),xishu_52.to(device)
            out13,out26,out52 = net(image)
            
            net.zero_grad()
            loss13 = loss_fun(out13,target13,traget_no_13,xishu_13)
            loss26 = loss_fun(out26,target26,traget_no_26,xishu_26)
            loss52 = loss_fun(out52,target52,traget_no_52,xishu_52)

            loss = loss13+loss26+loss52

            loss.backward()
            opt.step()
        now_best = loss
        print("loss:{}".format(loss))
        if now_best<history_best:
            history_best = now_best
            torch.save(net.state_dict(),weigh_path)
            print('保存成功')
        epoch +=1
        
